import os
import time
import re
from concurrent.futures import ThreadPoolExecutor
import requests
from loguru import logger
from tqdm import tqdm

# 基础配置
DEFAULT_DOWNLOAD_PATH = "downloads"
TEMP_DIR = "temp"
CHUNK_SIZE = 1024 * 1024  # 1MB
MAX_RETRY_TIMES = 3
MAX_CONCURRENT_DOWNLOADS = 5
RESUME_DOWNLOAD = True
MAX_FILENAME_LENGTH = 100  # 最大文件名长度

# 请求头设置
USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
REFERER = "https://www.bilibili.com"

def clean_filename(filename):
    """清理文件名，移除不合法字符并限制长度"""
    # 移除不合法字符
    filename = re.sub(r'[\\/:*?"<>|]', '', filename)
    # 移除多余的空格
    filename = ' '.join(filename.split())
    # 限制长度
    name, ext = os.path.splitext(filename)
    if len(name) > MAX_FILENAME_LENGTH:
        name = name[:MAX_FILENAME_LENGTH]
    return name + ext

class BilibiliDownloader:
    """哔哩哔哩视频下载器"""
    
    def __init__(self):
        """初始化下载器"""
        self.headers = {
            'User-Agent': USER_AGENT,
            'Referer': REFERER,
            'Origin': 'https://www.bilibili.com',
            'Accept': 'application/json, text/plain, */*',
            'Accept-Language': 'zh-CN,zh;q=0.9',
            'Accept-Encoding': 'gzip, deflate, br',
            'Connection': 'keep-alive'
        }
        self.session = requests.Session()
        self.session.headers.update(self.headers)
        self.quality_dict = {
            120: '4K',
            116: '1080P60',
            80: '1080P',
            74: '720P60',
            64: '720P',
            32: '480P',
            16: '360P'
        }
        
        # 创建必要的目录
        os.makedirs(DEFAULT_DOWNLOAD_PATH, exist_ok=True)
        os.makedirs(TEMP_DIR, exist_ok=True)
        
    def get_video_info(self, bvid):
        """获取视频信息"""
        try:
            url = f'https://api.bilibili.com/x/web-interface/view?bvid={bvid}'
            response = self.session.get(url)
            data = response.json()
            
            if data['code'] == 0:
                return data['data']
            else:
                logger.error(f"获取视频信息失败：{data['message']}")
                return None
        except Exception as e:
            logger.error(f"获取视频信息时发生错误：{str(e)}")
            return None
            
    def get_video_url(self, bvid, cid, quality=80):
        """获取视频下载地址"""
        try:
            url = f'https://api.bilibili.com/x/player/playurl?bvid={bvid}&cid={cid}&qn={quality}'
            response = self.session.get(url)
            data = response.json()
            
            if data['code'] == 0:
                return data['data']['durl'][0]['url']
            else:
                logger.error(f"获取视频下载地址失败：{data['message']}")
                return None
        except Exception as e:
            logger.error(f"获取视频下载地址时发生错误：{str(e)}")
            return None
            
    def download_video_part(self, url, start_byte, end_byte, temp_file):
        """下载视频片段"""
        headers = self.headers.copy()
        headers['Range'] = f'bytes={start_byte}-{end_byte}'
        
        for _ in range(MAX_RETRY_TIMES):
            try:
                response = self.session.get(url, headers=headers, stream=True)
                total_size = end_byte - start_byte + 1
                with open(temp_file, 'wb') as f:
                    with tqdm(total=total_size, unit='iB', unit_scale=True, desc=f"下载片段 {os.path.basename(temp_file)}") as pbar:
                        for chunk in response.iter_content(chunk_size=CHUNK_SIZE):
                            if chunk:
                                f.write(chunk)
                                pbar.update(len(chunk))
                return True
            except Exception as e:
                logger.error(f"下载视频片段时发生错误：{str(e)}")
                time.sleep(1)
        return False
        
    def merge_video_parts(self, temp_files, output_file):
        """合并视频片段"""
        try:
            with open(output_file, 'wb') as outfile:
                for temp_file in temp_files:
                    if os.path.exists(temp_file):
                        with open(temp_file, 'rb') as infile:
                            outfile.write(infile.read())
                        try:
                            os.remove(temp_file)
                        except:
                            pass
                            
            return True
        except Exception as e:
            logger.error(f"合并视频片段时发生错误：{str(e)}")
            return False
            
    def download_video(self, url, output_path, filename):
        """使用多线程下载视频"""
        try:
            if not os.path.exists(output_path):
                os.makedirs(output_path)
                
            # 清理文件名
            clean_name = clean_filename(filename)
            full_path = os.path.join(output_path, clean_name)
            temp_base = os.path.join(TEMP_DIR, os.path.splitext(clean_name)[0])
            
            # 获取文件大小
            response = self.session.head(url)
            file_size = int(response.headers.get('content-length', 0))
            
            # 如果启用断点续传，检查已下载的部分
            if RESUME_DOWNLOAD and os.path.exists(full_path):
                current_size = os.path.getsize(full_path)
                if current_size >= file_size:
                    logger.info(f"文件已存在且完整：{clean_name}")
                    return True
                    
            # 计算每个线程下载的字节范围
            chunk_size = file_size // MAX_CONCURRENT_DOWNLOADS
            ranges = []
            for i in range(MAX_CONCURRENT_DOWNLOADS):
                start_byte = i * chunk_size
                end_byte = start_byte + chunk_size - 1 if i < MAX_CONCURRENT_DOWNLOADS - 1 else file_size - 1
                ranges.append((start_byte, end_byte))
                
            # 创建临时文件
            temp_files = [f"{temp_base}.part{i}" for i in range(len(ranges))]
            
            print(f"\n开始下载：{clean_name}")
            print(f"文件大小：{file_size / 1024 / 1024:.2f}MB")
            print(f"分片数量：{len(ranges)}")
            
            # 使用线程池下载
            with ThreadPoolExecutor(max_workers=MAX_CONCURRENT_DOWNLOADS) as executor:
                futures = []
                for i, (start_byte, end_byte) in enumerate(ranges):
                    future = executor.submit(
                        self.download_video_part,
                        url,
                        start_byte,
                        end_byte,
                        temp_files[i]
                    )
                    futures.append(future)
                
                # 等待所有下载完成
                for future in futures:
                    future.result()
                        
            # 合并视频片段
            print("\n合并视频片段...")
            if not self.merge_video_parts(temp_files, full_path):
                return False
            
            logger.success(f"视频下载完成：{clean_name}")
            return True
            
        except Exception as e:
            logger.error(f"下载视频时发生错误：{str(e)}")
            return False
            
    def download_all_parts(self, bvid, output_path=DEFAULT_DOWNLOAD_PATH, quality=80):
        """下载所有分P"""
        video_info = self.get_video_info(bvid)
        if not video_info:
            return False
            
        # 获取所有分P信息
        parts = video_info.get('pages', [])
        if not parts:
            # 单P视频
            return self.download_by_bvid(bvid, output_path, quality)
            
        # 多P视频
        success = True
        for part in parts:
            cid = part['cid']
            part_name = f"{video_info['title']}_P{part['page']}_{part['part']}.mp4"
            
            # 获取视频地址
            video_url = self.get_video_url(bvid, cid, quality)
            if not video_url:
                success = False
                continue
                
            # 下载分P
            if not self.download_video(video_url, output_path, part_name):
                success = False
                
        return success
        
    def download_by_bvid(self, bvid, output_path=DEFAULT_DOWNLOAD_PATH, quality=80):
        """通过BV号下载视频"""
        # 获取视频信息
        video_info = self.get_video_info(bvid)
        if not video_info:
            return False
            
        # 获取视频下载地址
        video_url = self.get_video_url(bvid, video_info['cid'], quality)
        if not video_url:
            return False
            
        # 下载视频
        filename = f"{video_info['title']}.mp4"
        return self.download_video(video_url, output_path, filename)

def main():
    """主函数"""
    downloader = BilibiliDownloader()
    
    while True:
        print("\n=== B站视频下载工具 ===")
        print("1. 下载新视频")
        print("2. 退出")
        
        choice = input("\n请选择功能 (1-2): ")
        
        if choice == "1":
            bvid = input("请输入视频BV号：")
            print("\n可选视频质量：")
            for k, v in downloader.quality_dict.items():
                print(f"{k}: {v}")
            quality = int(input("请输入视频质量编号："))
            
            # 下载视频（自动处理分P）
            if downloader.download_all_parts(bvid, quality=quality):
                print("下载完成！")
            else:
                print("下载失败！")
                
        elif choice == "2":
            print("感谢使用！")
            break
        else:
            print("无效的选择，请重试！")

if __name__ == "__main__":
    main() 